"""
数据加载器类，用于从第二阶段加载和处理数据
"""

import json
import os
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path
from loguru import logger
from dataclasses import dataclass


@dataclass
class SessionData:
    """会话数据结构"""
    session_id: str
    metadata: Dict[str, Any]
    text_groups: List[Dict[str, Any]]
    images: List[Dict[str, Any]]
    tables: List[Dict[str, Any]]
    raw_data: List[Dict[str, Any]]


@dataclass
class ProcessingElement:
    """处理元素"""
    element_id: str
    content: Any
    category: str
    subcategory: str
    confidence: float
    metadata: Dict[str, Any]
    source_type: str  # 'text', 'image', 'table', 'raw'


class DataLoader:
    """数据加载器"""
    
    def __init__(self, storage_root: str = "E:/Program/gitee/slz/storage"):
        self.storage_root = Path(storage_root)
        
    def load_session_data(self, session_id: str) -> Optional[SessionData]:
        """加载会话数据"""
        session_path = self.storage_root / session_id
        
        if not session_path.exists():
            logger.error(f"会话目录不存在: {session_path}")
            return None
        
        try:
            # 加载元数据
            metadata = self._load_metadata(session_path)
            if not metadata:
                logger.error(f"无法加载会话元数据: {session_id}")
                return None
            
            # 加载各类数据
            text_groups = self._load_text_groups(session_path)
            images = self._load_images(session_path)
            tables = self._load_tables(session_path)
            raw_data = self._load_raw_data(session_path)
            
            return SessionData(
                session_id=session_id,
                metadata=metadata,
                text_groups=text_groups,
                images=images,
                tables=tables,
                raw_data=raw_data
            )
            
        except Exception as e:
            logger.error(f"加载会话数据失败 {session_id}: {e}")
            return None
    
    def _load_metadata(self, session_path: Path) -> Optional[Dict[str, Any]]:
        """加载元数据"""
        # 修复：实际元数据文件在metadata子目录下
        metadata_file = session_path / "metadata" / "session_metadata.json"
        
        if not metadata_file.exists():
            logger.warning(f"元数据文件不存在: {metadata_file}")
            return None
        
        try:
            with open(metadata_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载元数据失败: {e}")
            return None
    
    def _load_text_groups(self, session_path: Path) -> List[Dict[str, Any]]:
        """加载文本分组数据"""
        text_groups = []
        # 修复路径：实际数据在text/level3_groups目录下
        text_dir = session_path / "text" / "level3_groups"
        
        if not text_dir.exists():
            logger.warning(f"文本分组目录不存在: {text_dir}")
            return text_groups
        
        try:
            # 获取所有group文件
            group_files = list(text_dir.glob("group_*.json"))
            group_files.sort()  # 按文件名排序
            
            for group_file in group_files:
                try:
                    with open(group_file, 'r', encoding='utf-8') as f:
                        group_data = json.load(f)
                        text_groups.append(group_data)
                except Exception as e:
                    logger.error(f"加载文本分组文件失败 {group_file}: {e}")
            
            logger.info(f"成功加载 {len(text_groups)} 个文本分组")
            return text_groups
            
        except Exception as e:
            logger.error(f"加载文本分组数据失败: {e}")
            return text_groups
    
    def _load_images(self, session_path: Path) -> List[Dict[str, Any]]:
        """加载图像数据"""
        images = []
        images_dir = session_path / "images"
        
        if not images_dir.exists():
            logger.warning(f"图像目录不存在: {images_dir}")
            return images
        
        try:
            # 修复：加载实际的图像数据结构
            # 检查figures目录下的figure.json
            figures_file = images_dir / "figures" / "figure.json"
            if figures_file.exists():
                try:
                    with open(figures_file, 'r', encoding='utf-8') as f:
                        figures_data = json.load(f)
                        if isinstance(figures_data, list):
                            images.extend(figures_data)
                        else:
                            images.append(figures_data)
                except Exception as e:
                    logger.error(f"加载图像文件失败 {figures_file}: {e}")
            
            # 检查其他可能的图像子目录
            for subdir in ['charts', 'diagrams']:
                subdir_path = images_dir / subdir
                if subdir_path.exists():
                    for json_file in subdir_path.glob("*.json"):
                        try:
                            with open(json_file, 'r', encoding='utf-8') as f:
                                image_data = json.load(f)
                                if isinstance(image_data, list):
                                    images.extend(image_data)
                                else:
                                    images.append(image_data)
                        except Exception as e:
                            logger.error(f"加载图像文件失败 {json_file}: {e}")
            
            logger.info(f"成功加载 {len(images)} 个图像")
            return images
            
        except Exception as e:
            logger.error(f"加载图像数据失败: {e}")
            return images
    
    def _load_tables(self, session_path: Path) -> List[Dict[str, Any]]:
        """加载表格数据"""
        tables = []
        tables_dir = session_path / "tables"
        
        if not tables_dir.exists():
            logger.warning(f"表格目录不存在: {tables_dir}")
            return tables
        
        try:
            # 修复：加载实际的表格数据结构
            # 检查data目录下的general.json
            data_file = tables_dir / "data" / "general.json"
            if data_file.exists():
                try:
                    with open(data_file, 'r', encoding='utf-8') as f:
                        data_tables = json.load(f)
                        if isinstance(data_tables, list):
                            tables.extend(data_tables)
                        else:
                            tables.append(data_tables)
                except Exception as e:
                    logger.error(f"加载表格文件失败 {data_file}: {e}")
            
            # 检查structured目录下的structured.json
            structured_file = tables_dir / "structured" / "structured.json"
            if structured_file.exists():
                try:
                    with open(structured_file, 'r', encoding='utf-8') as f:
                        structured_tables = json.load(f)
                        if isinstance(structured_tables, list):
                            tables.extend(structured_tables)
                        else:
                            tables.append(structured_tables)
                except Exception as e:
                    logger.error(f"加载表格文件失败 {structured_file}: {e}")
            
            # 兼容性：也检查根目录下的JSON文件
            for table_file in tables_dir.glob("*.json"):
                try:
                    with open(table_file, 'r', encoding='utf-8') as f:
                        table_data = json.load(f)
                        if isinstance(table_data, list):
                            tables.extend(table_data)
                        else:
                            tables.append(table_data)
                except Exception as e:
                    logger.error(f"加载表格文件失败 {table_file}: {e}")
            
            logger.info(f"成功加载 {len(tables)} 个表格")
            return tables
            
        except Exception as e:
            logger.error(f"加载表格数据失败: {e}")
            return tables
    
    def _load_raw_data(self, session_path: Path) -> List[Dict[str, Any]]:
        """加载原始数据"""
        raw_data = []
        raw_dir = session_path / "raw"
        
        if not raw_dir.exists():
            logger.warning(f"原始数据目录不存在: {raw_dir}")
            return raw_data
        
        try:
            raw_files = list(raw_dir.glob("*.json"))
            
            for raw_file in raw_files:
                try:
                    with open(raw_file, 'r', encoding='utf-8') as f:
                        raw_item = json.load(f)
                        raw_data.append(raw_item)
                except Exception as e:
                    logger.error(f"加载原始数据文件失败 {raw_file}: {e}")
            
            logger.info(f"成功加载 {len(raw_data)} 个原始数据项")
            return raw_data
            
        except Exception as e:
            logger.error(f"加载原始数据失败: {e}")
            return raw_data
    
    def extract_processing_elements(self, session_data: SessionData) -> List[ProcessingElement]:
        """提取处理元素"""
        elements = []
        
        # 处理文本分组
        for group in session_data.text_groups:
            for element in group.get('elements', []):
                elements.append(ProcessingElement(
                    element_id=element.get('element_id', ''),
                    content=element.get('content', ''),
                    category=element.get('category', ''),
                    subcategory=element.get('subcategory', ''),
                    confidence=element.get('confidence', 0.0),
                    metadata=element.get('metadata', {}),
                    source_type='text'
                ))
        
        # 处理图像
        for image in session_data.images:
            elements.append(ProcessingElement(
                element_id=image.get('element_id', ''),
                content=image.get('content', ''),  # base64编码的图像内容
                category='image',
                subcategory=image.get('subcategory', ''),
                confidence=image.get('confidence', 1.0),
                metadata=image.get('metadata', {}),
                source_type='image'
            ))
        
        # 处理表格
        for table in session_data.tables:
            elements.append(ProcessingElement(
                element_id=table.get('element_id', ''),
                content=table.get('content', ''),
                category='table',
                subcategory=table.get('subcategory', ''),
                confidence=table.get('confidence', 1.0),
                metadata=table.get('metadata', {}),
                source_type='table'
            ))
        
        # 处理原始数据
        for raw_item in session_data.raw_data:
            elements.append(ProcessingElement(
                element_id=raw_item.get('id', ''),
                content=raw_item.get('content', ''),
                category='raw',
                subcategory=raw_item.get('type', ''),
                confidence=1.0,
                metadata=raw_item,
                source_type='raw'
            ))
        
        logger.info(f"提取了 {len(elements)} 个处理元素")
        return elements
    
    def get_latest_session_id(self) -> Optional[str]:
        """获取最新的会话ID"""
        try:
            session_dirs = [d for d in self.storage_root.iterdir() 
                          if d.is_dir() and d.name.startswith('session_')]
            
            if not session_dirs:
                logger.warning("未找到任何会话目录")
                return None
            
            # 按目录名排序，获取最新的
            session_dirs.sort(key=lambda x: x.name, reverse=True)
            latest_session = session_dirs[0].name
            
            logger.info(f"找到最新会话: {latest_session}")
            return latest_session
            
        except Exception as e:
            logger.error(f"获取最新会话ID失败: {e}")
            return None
    
    def filter_elements_by_type(
        self, 
        elements: List[ProcessingElement], 
        source_types: List[str]
    ) -> List[ProcessingElement]:
        """按类型过滤元素"""
        filtered = [e for e in elements if e.source_type in source_types]
        logger.info(f"按类型 {source_types} 过滤，得到 {len(filtered)} 个元素")
        return filtered
    
    def filter_elements_by_category(
        self, 
        elements: List[ProcessingElement], 
        categories: List[str]
    ) -> List[ProcessingElement]:
        """按类别过滤元素"""
        filtered = [e for e in elements if e.category in categories]
        logger.info(f"按类别 {categories} 过滤，得到 {len(filtered)} 个元素")
        return filtered
    
    def load_text_groups(self, session_id: str) -> List[Dict[str, Any]]:
        """加载文本分组数据 - 公共接口方法"""
        try:
            session_path = self.storage_root / session_id
            return self._load_text_groups(session_path)
        except Exception as e:
            logger.error(f"加载文本分组失败 {session_id}: {e}")
            return []
    
    def load_images(self, session_id: str) -> List[Dict[str, Any]]:
        """加载图像数据 - 公共接口方法"""
        try:
            session_path = self.storage_root / session_id
            return self._load_images(session_path)
        except Exception as e:
            logger.error(f"加载图像失败 {session_id}: {e}")
            return []
    
    def load_tables(self, session_id: str) -> List[Dict[str, Any]]:
        """加载表格数据 - 公共接口方法"""
        try:
            session_path = self.storage_root / session_id
            return self._load_tables(session_path)
        except Exception as e:
            logger.error(f"加载表格失败 {session_id}: {e}")
            return []
    
    def load_raw_data(self, session_id: str) -> List[Dict[str, Any]]:
        """加载原始数据 - 公共接口方法"""
        try:
            session_path = self.storage_root / session_id
            return self._load_raw_data(session_path)
        except Exception as e:
            logger.error(f"加载原始数据失败 {session_id}: {e}")
            return []

    def load_session_metadata(self, session_id: str) -> Dict[str, Any]:
        """加载会话元数据"""
        try:
            session_path = self.storage_root / session_id
            metadata_file = session_path / "metadata" / "session_metadata.json"
            
            if not metadata_file.exists():
                logger.warning(f"会话元数据文件不存在: {metadata_file}")
                return {
                    "session_id": session_id,
                    "created_at": "",
                    "status": "unknown",
                    "total_elements": 0,
                    "processing_time": 0
                }
            
            with open(metadata_file, 'r', encoding='utf-8') as f:
                metadata = json.load(f)
                
            logger.info(f"成功加载会话元数据: {session_id}")
            return metadata
            
        except Exception as e:
            logger.error(f"加载会话元数据失败: {e}")
            return {
                "session_id": session_id,
                "created_at": "",
                "status": "error",
                "total_elements": 0,
                "processing_time": 0
            }